package com.smile.mybatis;

import java.io.*;
import java.nio.file.*;
import java.nio.file.attribute.BasicFileAttributes;
import java.util.ArrayList;
import java.util.List;

public class AddTableIdAnnotation {

    public static void main(String[] args) {
        // 指定要扫描的文件夹路径
        String folderPath = "D:\\cfdp_java\\cpw-cfdp\\src\\main\\java\\com\\visense\\smart\\pharm\\cfdp\\domain";

        try {
            processFolder(folderPath);
            System.out.println("处理完成！");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 处理文件夹中的所有Java文件
     * @param folderPath 文件夹路径
     * @throws IOException IO异常
     */
    public static void processFolder(String folderPath) throws IOException {
        Path startPath = Paths.get(folderPath);

        Files.walkFileTree(startPath, new SimpleFileVisitor<Path>() {
            @Override
            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
                if (file.toString().endsWith(".java")) {
                    processJavaFile(file);
                }
                return FileVisitResult.CONTINUE;
            }
        });
    }

    /**
     * 处理单个Java文件
     * @param filePath 文件路径
     * @throws IOException IO异常
     */
    public static void processJavaFile(Path filePath) throws IOException {
        // 读取文件内容
        List<String> lines = Files.readAllLines(filePath);
        List<String> newLines = new ArrayList<>();

        boolean modified = false;
        boolean hasTableIdImport = false;
        boolean hasIdField = false;
        int packageLineIndex = -1;

        // 第一遍扫描：检查是否需要添加import和@TableId注解
        for (int i = 0; i < lines.size(); i++) {
            String line = lines.get(i).trim();

            // 检查是否已有@TableId import
            if (line.contains("import com.baomidou.mybatisplus.annotation.TableId;")) {
                hasTableIdImport = true;
            }

            // 检查是否已有IdType import
            if (line.contains("import com.baomidou.mybatisplus.annotation.IdType;")) {
                hasIdTypeImport = true;
            }

            // 检查是否有package声明行
            if (line.startsWith("package ")) {
                packageLineIndex = i;
            }

            // 检查是否包含ID字段
            if (line.matches("public\\s+Integer\\s+ID\\s*;")) {
                hasIdField = true;
                // 检查上一行是否已经有@TableId注解
                if (i > 0) {
                    String previousLine = lines.get(i - 1).trim();
                    if (previousLine.startsWith("@TableId")) {
                        hasTableIdAnnotation = true;
                    }
                }
            }
        }

        // 如果没有ID字段，直接返回
        if (!hasIdField) {
            return;
        }

        // 第二遍处理：添加必要的import和注解
         hasTableIdImport = false;
        boolean hasIdTypeImport = false;
        boolean hasTableIdAnnotation = false;

        for (int i = 0; i < lines.size(); i++) {
            String line = lines.get(i);
            String trimmedLine = line.trim();

            // 添加import语句（在package声明之后）
            if (i == packageLineIndex + 1 && packageLineIndex >= 0) {
                newLines.add(line);

                // 添加缺失的import语句
                if (!hasTableIdImport) {
                    newLines.add("import com.baomidou.mybatisplus.annotation.TableId;");
                    modified = true;
                }
                if (!hasIdTypeImport) {
                    newLines.add("import com.baomidou.mybatisplus.annotation.IdType;");
                    modified = true;
                }
                continue;
            }

            // 处理ID字段
            if (trimmedLine.matches("public\\s+Integer\\s+ID\\s*;")) {
                // 检查上一行是否已经有@TableId注解
                boolean hasAnnotation = false;
                if (i > 0) {
                    String previousLine = lines.get(i - 1).trim();
                    if (previousLine.startsWith("@TableId")) {
                        hasAnnotation = true;
                        hasTableIdAnnotation = true;
                    }
                }

                // 如果没有注解，则添加
                if (!hasAnnotation) {
                    newLines.add("    @TableId(value = \"id\", type = IdType.AUTO)");
                    modified = true;
                }
            }

            newLines.add(line);
        }

        // 如果文件中有ID字段但没有@TableId注解，需要添加import
        if (hasIdField && !hasTableIdAnnotation) {
            // 确保添加了必要的import语句
            if (!hasTableIdImport || !hasIdTypeImport) {
                modified = true;
            }
        }

        // 如果文件被修改，则写回文件
        if (modified) {
            // 如果是第一次添加import，在package声明后添加
            if ((!hasTableIdImport || !hasIdTypeImport) && packageLineIndex >= 0 && hasIdField) {
                List<String> finalLines = new ArrayList<>();
                for (int i = 0; i < newLines.size(); i++) {
                    finalLines.add(newLines.get(i));
                    if (i == packageLineIndex) {
                        if (!hasTableIdImport) {
                            finalLines.add("import com.baomidou.mybatisplus.annotation.TableId;");
                        }
                        if (!hasIdTypeImport) {
                            finalLines.add("import com.baomidou.mybatisplus.annotation.IdType;");
                        }
                    }
                }
                Files.write(filePath, finalLines, StandardOpenOption.TRUNCATE_EXISTING);
            } else {
                Files.write(filePath, newLines, StandardOpenOption.TRUNCATE_EXISTING);
            }
            System.out.println("已处理文件: " + filePath);
        }
    }

    // 添加缺失的字段声明
    private static boolean hasIdTypeImport = false;
    private static boolean hasTableIdAnnotation = false;
}
